import torch
from PIL import Image
from transformers import GenerationConfig
from auto_gptq.modeling import OvisGemma2GPTQForCausalLM
import pandas as pd
import numpy as np
import sys
import os
import time
import re

# 创建日志目录
log_dir = "logs"
os.makedirs(log_dir, exist_ok=True)

# 创建日志文件名（带时间戳）
timestamp = time.strftime("%Y%m%d_%H%M%S")
log_file = os.path.join(log_dir, f"ovis_EE_{timestamp}.log")

# 设置日志函数
def log_message(message, also_print=True):
    """将消息写入日志文件并可选择打印到控制台"""
    with open(log_file, "a", encoding="utf-8") as f:
        f.write(f"{message}\n")
    if also_print:
        print(message)

class EventMetrics:
    def __init__(self):
        self.event_types = [
            'Contact:Meet',
            'Justice:Arrest-Jail',
            'Conflict:Attack',
            'Conflict:Demonstrate',
            'Movement:Transport',
            'Contact:Phone-Write',
            'Life:Die',
            'Transaction:Transfer-Money'
        ]
        self.metrics = {event_type: {'TP': 0, 'FP': 0, 'FN': 0, 'total': 0} for event_type in self.event_types}
    
    def update(self, true_event, predicted_event):
        if true_event == 'none' or true_event in ['未识别', '显存不足', '图像加载失败', 'error']:
            return
            
        # 更新总数
        if true_event in self.metrics:
            self.metrics[true_event]['total'] += 1
            
        # 更新TP/FP/FN
        if true_event in self.metrics:
            if predicted_event == true_event:
                self.metrics[true_event]['TP'] += 1
            else:
                self.metrics[true_event]['FN'] += 1
                
        if predicted_event in self.metrics and predicted_event != true_event:
            self.metrics[predicted_event]['FP'] += 1
    
    def calculate_metrics(self):
        results = []
        for event_type in self.event_types:
            metrics = self.metrics[event_type]
            tp = metrics['TP']
            fp = metrics['FP']
            fn = metrics['FN']
            total = metrics['total']
            
            precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
            recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
            f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0
            
            results.append({
                'event_type': event_type,
                'precision': precision,
                'recall': recall,
                'f1': f1,
                'support': total
            })
        return results
    
    def print_metrics(self):
        results = self.calculate_metrics()
        
        # 打印表头
        log_message("")
        log_message("-" * 80)
        log_message(f"{'事件类型':<30} {'精确度':>10} {'召回率':>10} {'F1值':>10} {'支持数':>10}")
        log_message("-" * 80)
        
        # 打印每种事件类型的指标
        for result in results:
            log_message(f"{result['event_type']:<30} "
                     f"{result['precision']:>10.4f} "
                     f"{result['recall']:>10.4f} "
                     f"{result['f1']:>10.4f} "
                     f"{result['support']:>10}")
        
        log_message("-" * 80)

# 数据路径
data_path = r"/mnt/6t/lyh/1/M2E2/M2E2/text/mee.tsv"
data = pd.read_csv(data_path, sep='\t')
        
def load_image(image_path):
    try:
        image = Image.open(image_path)
        return image
    except Exception as e:
        log_message(f"无法加载图片 {image_path}: {e}")
        return None

# 加载模型
load_device = "cuda:0"
model_path = "/mnt/6t/lyh/1/Ovis1.6-Gemma2-9B-GPTQ-Int4"
model = OvisGemma2GPTQForCausalLM.from_quantized(
    model_path,
    device=load_device,
    trust_remote_code=True,
    low_cpu_mem_usage=True
)
model.model.generation_config = GenerationConfig.from_pretrained(model_path)
text_tokenizer = model.get_text_tokenizer()
visual_tokenizer = model.get_visual_tokenizer()

def extract_quoted_text(text):
    """
    从文本中提取单引号或双引号包裹的字符串，并消除空格
    如果没有找到引号包裹的内容，则返回原始文本去除空格后的结果
    """
    # 匹配单引号或双引号中的内容
    single_quote_pattern = r"'([^']*)'|'([^']*)'|'([^']*)'"  # 包括不同类型的单引号
    double_quote_pattern = r'"([^"]*)"'
    
    # 先尝试匹配单引号
    single_quotes = re.findall(single_quote_pattern, text)
    if single_quotes:
        # re.findall对于多个捕获组会返回元组列表，需要找到非空的那个
        for quote_tuple in single_quotes:
            for quote in quote_tuple:
                if quote:
                    # 删除所有空格
                    return quote.replace(" ", "")
    
    # 再尝试匹配双引号
    double_quotes = re.findall(double_quote_pattern, text)
    if double_quotes:
        # 删除所有空格
        return double_quotes[0].replace(" ", "")
    
    # 如果没有找到引号包裹的内容，则返回原始文本去除所有空格
    return text.replace(" ", "")

def normalize_event_type(event_text):
    """
    规范化事件类型文本
    1. 提取引号中的内容
    2. 去除空格
    3. 检查是否匹配允许的事件类型
    """
    # 提取引号中的内容并去除空格
    extracted_text = extract_quoted_text(event_text)
    
    # 定义允许的事件类型（无空格版本）
    allowed_events_no_space = {
        'Justice:Arrest-Jail': 'Justice:Arrest-Jail',
        'Conflict:Attack': 'Conflict:Attack',
        'Movement:Transport': 'Movement:Transport',
        'Contact:Phone-Write': 'Contact:Phone-Write',
        'Life:Die': 'Life:Die',
        'Conflict:Demonstrate': 'Conflict:Demonstrate',
        'Transaction:Transfer-Money': 'Transaction:Transfer-Money',
        'Contact:Meet': 'Contact:Meet',
        'none': 'none'
    }
    
    # 创建无空格版本的映射
    allowed_events_map = {key.replace(" ", ""): value for key, value in allowed_events_no_space.items()}
    
    # 检查是否完全匹配允许的事件类型（无空格比较）
    if extracted_text.lower() in [event.lower().replace(" ", "") for event in allowed_events_no_space.keys()]:
        for no_space_key, original_key in allowed_events_map.items():
            if extracted_text.lower() == no_space_key.lower():
                return original_key
    
    # 如果没有完全匹配，则检查是否包含允许的事件类型
    for no_space_key, original_key in allowed_events_map.items():
        if no_space_key.lower() in extracted_text.lower():
            return original_key
    
    # 如果都不匹配，返回提取的文本
    return extracted_text

def event_extraction_predict(text, image_path):
    image = load_image(image_path)
    if image is None:
        return "error"

    # 事件抽取提示词
    prompt = (
        "请根据下方文本和图像内容进行事件抽取，抽取事件类型可以参考以下事件类型，也可以选择其他事件类型："
        "'Justice:Arrest-Jail','Conflict:Attack','Movement:Transport','Contact:Phone-Write',"
        "'Life:Die','Conflict:Demonstrate','Transaction:Transfer-Money','Contact:Meet','none'。"
        "文本和图片证据不足的情况下直接输出none。输出内容只能包含事件类型，不要添加其他内容，不要添加其他符号，以及和事件类型无关的任何内容。"
    )
    query = f'<image>\n{prompt}\n{text}'

    # 格式化输入
    prompt, input_ids, pixel_values = model.preprocess_inputs(query, [image])
    attention_mask = torch.ne(input_ids, text_tokenizer.pad_token_id)
    input_ids = input_ids.unsqueeze(0).to(device=model.device)
    attention_mask = attention_mask.unsqueeze(0).to(device=model.device)
    pixel_values = [pixel_values.to(dtype=visual_tokenizer.dtype, device=visual_tokenizer.device)]

    # 生成输出
    with torch.inference_mode():
        gen_kwargs = dict(
            max_new_tokens=1024,
            do_sample=False,
            top_p=None,
            top_k=None,
            temperature=None,
            repetition_penalty=None,
            eos_token_id=model.generation_config.eos_token_id,
            pad_token_id=text_tokenizer.pad_token_id,
            use_cache=True
        )
        try:
            output_ids = model.generate(input_ids, pixel_values=pixel_values, attention_mask=attention_mask, **gen_kwargs)[0]
            output = text_tokenizer.decode(output_ids, skip_special_tokens=True)
            
            # 提取模型输出中的事件类型
            raw_prediction = output.split(query)[-1].strip()
            return raw_prediction
            
        except torch.cuda.OutOfMemoryError:
            log_message(f"显存不足，跳过当前数据：{image_path}")
            torch.cuda.empty_cache()
            return "error"

# 初始化计数器
TP = 0
FP = 0
FN = 0

# 初始化事件类型指标统计
event_metrics = EventMetrics()

log_message("开始事件抽取评估...")
for i, row in data.iterrows():
    log_message(f"--- 示例 {i+1}/{len(data)} ---")
    image_path = f"/mnt/6t/lyh/1/M2E2/M2E2/image/image/image1/{row['#1 ImageID']}"
    
    # 检查图片是否存在
    if not os.path.exists(image_path):
        log_message(f"图片不存在: {image_path}")
        continue
    
    # 获取真实事件类型
    true_event = row['#3 String']
    
    # 预测事件类型
    raw_prediction = event_extraction_predict(row['#2 String'], image_path)
    log_message(f"模型原始输出: {raw_prediction}")
    # 提取引号中的内容并规范化
    predicted_event = normalize_event_type(raw_prediction)
    
    # 更新事件类型指标
    event_metrics.update(true_event, predicted_event)
    
    # 简化输出，只保留关键信息
    log_message(f"真实事件类型: {true_event}")
    log_message(f"预测事件类型: {predicted_event}")
    
    # 评估结果
    result_type = ""
    
    # 按照要求计算TP、FP、FN
    # 1. 如果预测值与真实值相等为TP
    if predicted_event == true_event:
        TP += 1
        result_type = "TP"
    # 2. 如果预测值不为none且与真实值不相等为FP
    elif predicted_event != 'none' and predicted_event != 'error':
        FP += 1
        result_type = "FP"
    
    # 3. 如果预测值为none且真实值不为none为FN
    if predicted_event == 'none' and true_event != 'none':
        FN += 1
        result_type = "FN"
    
    log_message(f"结果类型: {result_type}")
    log_message(f"当前统计 - TP: {TP}, FP: {FP}, FN: {FN}")
    
    # 每10个样本计算一次中间结果
    if (i + 1) % 10 == 0 or i == len(data) - 1:
        precision = TP / (TP + FP) if TP + FP > 0 else 0
        recall = TP / (TP + FN) if TP + FN > 0 else 0
        f1_score = (2 * precision * recall) / (precision + recall) if precision + recall > 0 else 0
        
        log_message(f"处理样本数: {i+1}/{len(data)}")
        log_message(f"TP: {TP}, FP: {FP}, FN: {FN}")
        log_message(f"精确度: {precision:.4f}, 召回率: {recall:.4f}, F1值: {f1_score:.4f}")

# 计算总体精确度、召回率和F1值
precision = TP / (TP + FP) if TP + FP > 0 else 0
recall = TP / (TP + FN) if TP + FN > 0 else 0
f1_score = (2 * precision * recall) / (precision + recall) if precision + recall > 0 else 0

# 输出总体结果
log_message("\n=== 总体结果统计 ===")
log_message(f"TP: {TP}, FP: {FP}, FN: {FN}")
log_message(f"精确度: {precision:.4f}, 召回率: {recall:.4f}, F1值: {f1_score:.4f}")

# 输出各事件类型的指标统计
log_message("\n=== 各事件类型指标统计 ===")
event_metrics.print_metrics()